# Tool to prepare lidar data from Oxford RobotCar dataset.
# Licensed under the Apache License
import random
import time
import argparse, tqdm, os, errno
import numpy as np
from transform import (
    build_se3_transform, 
    se3_transform, 
    inverse_transform,
    compose_transform,
    frame_transform
)
from radar import load_radar, radar_polar_to_cartesian
import cv2
import matplotlib.pyplot as plt
import PIL.Image as Img
import PIL
LIDAR_PATH = None
LIDAR_HISTORY_PATH = None
RADAR_EXTRINSICS = [-0.71813, 0.12, -0.54479, 0, 0.05, 0]
LEFT_LIDAR_EXTRINSICS = [-0.60072, -0.34077, -0.26837, -0.0053948, -0.041998, -3.1337]
RIGHT_LIDAR_EXTRINSICS = [-0.61153, 0.55676, -0.27023, 0.0027052, -0.041999, -3.1357]
RADAR_ROT, RADAR_POS = se3_transform(RADAR_EXTRINSICS)
RADAR_INV_ROT, RADAR_INV_POS = inverse_transform(RADAR_ROT, RADAR_POS)
LEFT_LIDAR_ROT, LEFT_LIDAR_POS = se3_transform(LEFT_LIDAR_EXTRINSICS)
LEFT_LIDAR_INV_ROT, LEFT_LIDAR_INV_POS = inverse_transform(LEFT_LIDAR_ROT, LEFT_LIDAR_POS)
RIGHT_LIDAR_ROT, RIGHT_LIDAR_POS = se3_transform(RIGHT_LIDAR_EXTRINSICS)
RIGHT_LIDAR_INV_ROT, RIGHT_LIDAR_INV_POS = inverse_transform(RIGHT_LIDAR_ROT, RIGHT_LIDAR_POS)

def tensor_to_image(tensor):
        tensor = tensor*255
        tensor = np.array(tensor, dtype=np.uint8)
        if np.ndim(tensor)>3:
            assert tensor.shape[0] == 1
            tensor = tensor[0]
        return Img.fromarray(tensor)

def visualize_LiDAR_PointCloud(lidar_data):
    lidar_location = lidar_data[0:3,:]
    lidar_intensity = lidar_data[3:,:]
    lidar_intensity = lidar_intensity % 255.0
    arr = np.squeeze(np.asarray(lidar_location))
    img_loc = Img.fromarray(arr).convert("L")
    # img_loc.save("/home/aero/mvdnet/lidar_intensity.jpeg")
    # img_loc.show()
    # print(type(lidar_intensity))
    print(arr)
    # print(lidar_location)
    # lidar_intensity = lidar_data[3:,:]
    # img_loc.show()
    pass

def lidar_pc2pixor(lidar_data, delta_l=0.2, pixel_l=320, h1=-1.0, h2=2.5, delta_h=0.1):
    l1 = (-pixel_l/2) * delta_l
    l2 = (pixel_l/2) * delta_l
    pixel_h = np.int(np.round((h2 - h1) / delta_h))

    lidar_data[:,0] = -lidar_data[:,0]
    idx_x = np.logical_and(lidar_data[:,0] >= l1, lidar_data[:,0] < l2)
    idx_y = np.logical_and(lidar_data[:,1] >= l1, lidar_data[:,1] < l2)
    idx_z = np.logical_and(lidar_data[:,2] >= h1, lidar_data[:,2] < h2)
    idx_valid = np.logical_and(idx_z, np.logical_and(idx_y, idx_x))
    lidar_data = lidar_data[idx_valid, :]

    lidar_bev_idx = np.zeros([len(lidar_data), 2])
    lidar_bev_idx[:,0] = np.floor((lidar_data[:,0] - l1) / delta_l)
    lidar_bev_idx[:,1] = np.floor((lidar_data[:,1] - l1) / delta_l)
    lidar_bev_idx[lidar_bev_idx == pixel_l] = pixel_l - 1
    lidar_bev_idx = lidar_bev_idx.astype(np.int)

    lidar_height_idx = np.floor((lidar_data[:,2] - h1) / delta_h)
    lidar_height_idx[lidar_height_idx == pixel_h] = pixel_h - 1
    lidar_height_idx = lidar_height_idx.astype(np.int)

    lidar_intensity = np.zeros([pixel_l, pixel_l])
    lidar_occupancy = np.zeros([pixel_l, pixel_l, pixel_h])
    for i in range(len(lidar_bev_idx)):
        lidar_occupancy[lidar_bev_idx[i,0], lidar_bev_idx[i,1], lidar_height_idx[i]] = 1
        lidar_intensity[lidar_bev_idx[i,0], lidar_bev_idx[i,1]] = max(lidar_data[i,3], \
            lidar_intensity[lidar_bev_idx[i,0], lidar_bev_idx[i,1]])

    return lidar_intensity, lidar_occupancy


def main(args, HISTORY_NUM=4, OFFSET=0,FRAME_RATIO=5):
    lidar_odometry_path = os.path.join(args.data_path, 'vo/vo.csv')
    lidar_odometry = np.genfromtxt(lidar_odometry_path, delimiter=',')[1:]
    lidar_odometry_rot = []
    lidar_odometry_pos = []
    lidar_odometry_timestamp = []
    for sample in lidar_odometry:
        sample_rot, sample_pos = se3_transform(sample[2:])
        lidar_odometry_rot.append(sample_rot)
        lidar_odometry_pos.append(sample_pos)
        lidar_odometry_timestamp.append(int(sample[1]))
    lidar_odometry_timestamp.append(sample[0])
    lidar_odometry = {'rot':lidar_odometry_rot, 
                    'pos':lidar_odometry_pos,
                    'timestamp':lidar_odometry_timestamp}

    radar_timestamp_path = os.path.join(args.data_path, 'radar.timestamps')
    radar_timestamp = np.loadtxt(radar_timestamp_path)[:,0]

    left_lidar_timestamp_path = os.path.join(args.data_path, 'velodyne_left.timestamps')
    left_lidar_timestamp = np.loadtxt(left_lidar_timestamp_path)[:,0]

    right_lidar_timestamp_path = os.path.join(args.data_path, 'velodyne_right.timestamps')
    right_lidar_timestamp = np.loadtxt(right_lidar_timestamp_path)[:,0]

    radar_odometry_path = os.path.join(args.data_path, 'gt/radar_odometry.csv')
    radar_odometry = np.genfromtxt(radar_odometry_path, delimiter=',')[1:,np.r_[9,8,2:8]]

    num_history = HISTORY_NUM
    num_frame = len(radar_timestamp)-1
    lidar_radar_frame_ratio = FRAME_RATIO
    right_lidar_timestamp_i = 0
    # proc_time = 0 # total time of data preparation
    # proc_cnt = 0
    # start_time = time.time_ns()
    # for radar_timestamp_i in tqdm.tqdm(range(len(radar_timestamp) - 1)):
    for radar_timestamp_i in tqdm.tqdm(range(2)):
        while right_lidar_timestamp[right_lidar_timestamp_i] <= radar_timestamp[radar_timestamp_i]:
            right_lidar_timestamp_i += 1
        right_lidar_timestamp_i += OFFSET
        lidar_data_all = []
        for right_lidar_timestamp_j in range(right_lidar_timestamp_i-1, right_lidar_timestamp_i+lidar_radar_frame_ratio+1):
            right_lidar_filename = os.path.join(args.data_path, 'velodyne_right', str(int(right_lidar_timestamp[right_lidar_timestamp_j])) + '.bin')
            right_lidar_data = np.fromfile(right_lidar_filename, dtype=np.float32)
            right_lidar_data = np.matrix(np.reshape(right_lidar_data, (4, -1)))
            
            left_lidar_timestamp_closest = left_lidar_timestamp[min(range(len(left_lidar_timestamp)), key=lambda ii: abs(left_lidar_timestamp[ii] - right_lidar_timestamp[right_lidar_timestamp_j]))]
            left_lidar_filename = os.path.join(args.data_path, 'velodyne_left', str(int(left_lidar_timestamp_closest)) + '.bin')
            left_lidar_data = np.fromfile(left_lidar_filename, dtype=np.float32)
            left_lidar_data = np.matrix(np.reshape(left_lidar_data, (4, -1)))

            frame_rot, frame_pos = frame_transform(left_lidar_timestamp_closest, right_lidar_timestamp[right_lidar_timestamp_j], lidar_odometry)
            frame_rot, frame_pos = compose_transform(LEFT_LIDAR_ROT, LEFT_LIDAR_POS, frame_rot, frame_pos)
            frame_rot, frame_pos = compose_transform(frame_rot, frame_pos, RIGHT_LIDAR_INV_ROT, RIGHT_LIDAR_INV_POS)
            left_lidar_data[0:3,:] = np.matrix(frame_rot.as_dcm()) * left_lidar_data[0:3,:] + np.tile(np.matrix(frame_pos).T, (1, left_lidar_data.shape[1]))

            lidar_data = np.concatenate((right_lidar_data, left_lidar_data), axis=1)
            lidar_self = np.logical_and(np.logical_and(lidar_data[0,:] > -2.1, lidar_data[0,:] < 2), np.logical_and(lidar_data[1,:] > -0.5, lidar_data[1,:] < 1.4))
            lidar_sel = np.where(np.logical_not(lidar_self))[1]
            
            lidar_data = lidar_data[:,lidar_sel]
            lidar_location = lidar_data[0:3,:]
            lidar_intensity = lidar_data[3:,:]
            lidar_data = np.concatenate((lidar_location, lidar_intensity), axis=0)
            # visualize_LiDAR_PointCloud(lidar_data)
            frame_rot, frame_pos = frame_transform(right_lidar_timestamp[right_lidar_timestamp_j], radar_timestamp[radar_timestamp_i], lidar_odometry)
            frame_rot, frame_pos = compose_transform(RIGHT_LIDAR_ROT, RIGHT_LIDAR_POS, frame_rot, frame_pos)
            frame_rot, frame_pos = compose_transform(frame_rot, frame_pos, RADAR_INV_ROT, RADAR_INV_POS)
            frame_rot = np.matrix(frame_rot.as_dcm())
            frame_pos = np.matrix(frame_pos).T
            frame_pos = np.tile(frame_pos, (1, lidar_location.shape[1]))

            lidar_location = np.array((frame_rot * np.matrix(lidar_location) + frame_pos))
            print("^"*50)
            # print("np.matrix(lidar_location)")
            # print(np.matrix(lidar_location))
            # print("lidar_location")
            # print(lidar_location.shape[0])
            # print(lidar_location.shape[1])
            # print(lidar_location)
            lidar_angle = np.arctan2(lidar_location[1,:], lidar_location[0,:])
            

            lidar_intensity = lidar_intensity / 255.0
            # print("V"*50)
            if right_lidar_timestamp_j == right_lidar_timestamp_i - 1:
                sector_bound = 0
                sector_width = 2*np.pi/(lidar_radar_frame_ratio+1)
            elif right_lidar_timestamp_j == right_lidar_timestamp_i + lidar_radar_frame_ratio:
                sector_bound = lidar_radar_frame_ratio*2*np.pi/(lidar_radar_frame_ratio+1)
                sector_width = 2*np.pi/(lidar_radar_frame_ratio+1)
            else:
                sector_bound = (right_lidar_timestamp_j-right_lidar_timestamp_i)*2*np.pi/(lidar_radar_frame_ratio+1)
                sector_width = 2*2*np.pi/(lidar_radar_frame_ratio+1)
            print("i")
            print(right_lidar_timestamp_j-right_lidar_timestamp_i)
            print("sector_width")
            print(sector_width* 180.0 / np.pi )
            print("sector_bound")
            print(sector_bound* 180.0 / np.pi )
            print("lidar_angle0")
            # print(lidar_angle* 180.0 / np.pi )
            print(np.amax(lidar_angle* 180.0 / np.pi))
            print(np.amin(lidar_angle* 180.0 / np.pi))
            lidar_angle = np.mod(lidar_angle - sector_bound, 2*np.pi)
            print("lidar_angle1")
            # print(lidar_angle* 180.0 / np.pi )
            print(np.amax(lidar_angle* 180.0 / np.pi))
            print(np.amin(lidar_angle* 180.0 / np.pi))
            print("lidar_sel")
            print(lidar_sel)
            lidar_sel = np.logical_and(lidar_angle >= 0, lidar_angle < sector_width)
            print("lidar_sel")
            print(lidar_sel)
            print("V"*50)
            lidar_location = lidar_location[:,lidar_sel]
            # print(lidar_location.shape[0])
            # print(lidar_location.shape[1])
            # print("-"*50)
            lidar_intensity = lidar_intensity[:,lidar_sel]
            lidar_data = np.concatenate((lidar_location, lidar_intensity), axis=0)
            lidar_data_all.append(lidar_data)

        lidar_data_all = np.concatenate(lidar_data_all, axis=1).T
        test=np.array(lidar_data_all, dtype=np.float32)
        test=test.reshape((-1, 4))
        print("test.shape0")
        print(test.shape[0])
        print(test.shape[1])
        lidar_intensity, lidar_occupancy = lidar_pc2pixor(lidar_data=test)
        lidar_intensity=np.ascontiguousarray(np.expand_dims(lidar_intensity, -1).transpose(2, 0, 1))
        # print(lidar_intensity.shape)
        # print(lidar_intensity[0])
        # print(np.amax(lidar_intensity[0]))
        lidar_intensity*=255
        img=np.array(lidar_intensity,dtype=np.uint8)
        print(img.shape)
        # print(np.amax(img))
        img_loc = Img.fromarray(img[0]).convert("L")
        img_loc.save(f"R{FRAME_RATIO}_{radar_timestamp_i}.jpeg")
        # print(type(lidar_intensity))
        # print(test.shape[1])
        # plt.imshow(img[0].view(320,320))
        # lidar_data_path = os.path.join(LIDAR_PATH, str(int(radar_timestamp[radar_timestamp_i])) + '.bin')
        # lidar_data_all.astype(np.float32).tofile(lidar_data_path)

        # if radar_timestamp_i >= num_history:
        #     invT = np.eye(4)
        #     for radar_timestamp_j in range(1, num_history+1):
        #         frameXYZRPY = np.array([radar_odometry[radar_timestamp_i-radar_timestamp_j,2], radar_odometry[radar_timestamp_i-radar_timestamp_j,3], 0, 0, 0, radar_odometry[radar_timestamp_i-radar_timestamp_j,7]])
        #         frameT = build_se3_transform(frameXYZRPY)
        #         invT = frameT * invT
        #         T = np.linalg.inv(invT).T
        #         lidar_T_path = os.path.join(LIDAR_HISTORY_PATH, str(int(radar_timestamp[radar_timestamp_i])) + '_' + str(radar_timestamp_j) + '_T.bin')
        #         T.astype(np.float32).tofile(lidar_T_path)
        #         lidar_history_dst_path = os.path.join(LIDAR_HISTORY_PATH, str(int(radar_timestamp[radar_timestamp_i])) + '_' + str(radar_timestamp_j) + '.bin')
        #         lidar_history_src_path = os.path.join(LIDAR_PATH, str(int(radar_timestamp[radar_timestamp_i-radar_timestamp_j])) + '.bin')
        #         try:
        #             os.symlink(lidar_history_src_path, lidar_history_dst_path)
        #         except OSError as e:
        #             if e.errno == errno.EEXIST:
        #                 os.remove(lidar_history_dst_path)
        #                 os.symlink(lidar_history_src_path, lidar_history_dst_path)
        #             else:
        #                 raise e
    #     proc_cnt += 1
    # end_time = time.time_ns()
    # print(f"proc_time avg is {proc_time / (proc_cnt*1000000) }")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Prepare lidar data from Oxford RobotCar radar dataset')
    parser.add_argument('--data_path', type=str, required=True, help='path to the data record folder')
    parser.add_argument('--historyNum', type=int, default=4, help='history num')
    parser.add_argument('--offset', type=int, default=0, help='sensor offset')
    parser.add_argument('--frame_ratio', type=int, default=5, help='lidar radar frame ratio')
    args = parser.parse_args()
    HISTORY_NUM = args.historyNum
    OFFSET = args.offset
    FRAME_RATIO = args.frame_ratio
    # processed_path = os.path.join(args.data_path, f'num{HISTORY_NUM}/eval_processed')
    # renamed_processed_path = os.path.join(args.data_path, f'num{HISTORY_NUM}/eval_processed{str(OFFSET)}')
    # offset_indicator_path = os.path.join(args.data_path, f'num{HISTORY_NUM}/eval_processed/offset{str(OFFSET)}')
    # if not os.path.isdir(processed_path):
    #     os.mkdir(processed_path)
    # if not os.path.isdir(offset_indicator_path):
    #     os.mkdir(offset_indicator_path)
    # if FRAME_RATIO != 5:
    #     ratio_indicator_path = os.path.join(args.data_path, f'num{HISTORY_NUM}/eval_processed/ratio{str(FRAME_RATIO)}')
    #     if not os.path.isdir(ratio_indicator_path):
    #         os.mkdir(ratio_indicator_path)
    #     addPlusMinusToNum = lambda i: ("+" if i > 0 else "") + str(i)
    #     renamed_processed_path = os.path.join(args.data_path, f'num{HISTORY_NUM}/eval_processed_r{str(FRAME_RATIO)}{str(addPlusMinusToNum(OFFSET))}')
    # LIDAR_PATH = os.path.join(processed_path, 'lidar')
    # if not os.path.isdir(LIDAR_PATH):
    #     os.mkdir(LIDAR_PATH)

    # LIDAR_HISTORY_PATH = os.path.join(processed_path, 'lidar_history')
    # if not os.path.isdir(LIDAR_HISTORY_PATH):
    #     os.mkdir(LIDAR_HISTORY_PATH)
    main(args,HISTORY_NUM,OFFSET,FRAME_RATIO)
    # try:
    #     main(args,HISTORY_NUM,OFFSET,FRAME_RATIO)
    # except:
    #     print('Not all added')

    # if os.path.isdir(renamed_processed_path):
    #     renamed_processed_path = os.path.join(args.data_path, f'num{HISTORY_NUM}/eval_processed{str(OFFSET)}-copy--{str(random.randint(1,10))}')
    #     os.rename(processed_path,renamed_processed_path)
    #     raise SystemExit(f'Error: Existing Folder processed{str(OFFSET)} ')
    # os.rename(processed_path,renamed_processed_path)